from collections import Callable
import os
import h5py
import numpy as np
from dataset import map


class _read_partly_from_h5_file(Callable):
    def __init__(self, h5_file_path):
        assert os.path.isfile(h5_file_path)
        self._h5_file_path = h5_file_path

    def __call__(self, path_starts_shape):
        path, starts, shape = path_starts_shape
        slices = tuple(slice(max(st, 0), st + sp) for st, sp in zip(starts, shape))
        with h5py.File(self._h5_file_path, 'r') as f:
            raw = f[path][slices]
            mshape = f[path].shape
        pad_width = np.asarray([(0 - b, b + s - ms) for b, s, ms in zip(starts, shape, mshape)], 'int')
        pad_width[pad_width < 0] = 0
        paded = np.pad(raw, pad_width, 'edge')
        assert paded.shape == tuple(shape)
        return paded


def read_partly_from_h5_file(dataset, h5_file_path):
    return map(_read_partly_from_h5_file(h5_file_path), dataset)
